/*
 * HXD1D register allocation: using graph coloring.
 */
#include "HRV64Generator.h"
#include "HExceptions.h"

/*
 * Registers that can be allocated using graph coloring.
 * At least 3 GFRs and 3 FPRs are needed to
 * support spilled variables. So, `$s9`, `$s10`, `$s11`,
 * `$fs9`, `$fs10`, and `$fs11` are selected.
 */
const RV64GFR HWorkingIGraph::Available_GPR[HWorkingIGraph::NAGPR] = {
	RV64GFR::T0, RV64GFR::T1, RV64GFR::T2, RV64GFR::T3, RV64GFR::T4,
	RV64GFR::T5, RV64GFR::T6, RV64GFR::S0, RV64GFR::S1, RV64GFR::S2,
	RV64GFR::S3, RV64GFR::S4, RV64GFR::S5, RV64GFR::S6, RV64GFR::S7,
	RV64GFR::S8, //RV64GFR::S9, RV64GFR::S10, RV64GFR::S11,
	RV64GFR::A7, RV64GFR::A6, RV64GFR::A5, RV64GFR::A4, RV64GFR::A3,
	RV64GFR::A2, RV64GFR::A1, RV64GFR::A0,
};
const RV64GFR HWorkingIGraph::Available_FPR[HWorkingIGraph::NAFPR] = {
	RV64GFR::FT0, RV64GFR::FT1, RV64GFR::FT2, RV64GFR::FT3, RV64GFR::FT4, RV64GFR::FT5,
	RV64GFR::FT6, RV64GFR::FT7, RV64GFR::FT8, RV64GFR::FT9, RV64GFR::FT10, RV64GFR::FT11,
	RV64GFR::FS0, RV64GFR::FS1, RV64GFR::FS2, RV64GFR::FS3, RV64GFR::FS4, RV64GFR::FS5,
	RV64GFR::FS6, RV64GFR::FS7, RV64GFR::FS8, //RV64GFR::FS9, RV64GFR::FS10, RV64GFR::FS11,
	RV64GFR::FA7, RV64GFR::FA6, RV64GFR::FA5, RV64GFR::FA4, RV64GFR::FA3, RV64GFR::FA2,
	RV64GFR::FA1, RV64GFR::FA0
};

extern std::ofstream log_output;

/*
 * Try to simplify the graph by removing a low degree
 * vertex and push it into the stack.
 * Return non-zero integer on success.
 */
uint64_t HWorkingIGraph::simplify(void)
{
	size_t i, n = vertexes_.size();
	std::vector<HIGVertex>::iterator vit;
	for (i = 0ul, vit = vertexes_.begin(); i < n; ++i, ++vit)
		if (vit->valid && vit->allocated == false	// valid and not preallocated
			&& search_copy(i) == SIZE_MAX	// has no copy edge
			&& degree_of(i) < (vit->is_float ? NAFPR : NAGPR))	// low degree
			break;
	if (i >= n)
		return 0ul;
	else
	{
		/*
		 * If a low degree vertex that has no copy edge is found,
		 * push it into the stack and delete it.
		 * But its `adjlist` info must be kept!
		 */
		simplified_stack_.push_back(i);
		vit->valid = false;
		for (auto u : vit->adjlist)
			vertexes_.at(u).adjlist.erase(i);
		return 1ul;
	}
}

/*
 * Try to coalesce two vertexes using Briggs strategy.
 * Return non-zero integer on success.
 */
uint64_t HWorkingIGraph::coalesce(void)
{
	// Adjacencies after coalescence
	std::set<size_t> nadj;
	// Number of high degree adjacencies
	size_t hd_ymr;
	// Flag of success
	uint64_t flag_found = 0ul;
	// iterator
	std::vector<HIGCopy>::iterator cpit, cpit2;

	for (cpit = copies_.begin(); cpit != copies_.end(); ++cpit)
		if (cpit->status == HIGCopy::NORMAL)
		{
			const HIGVertex & v1 = vertexes_.at(cpit->igv1),
				& v2 = vertexes_.at(cpit->igv2);
			nadj.clear();
			std::set_union(v1.adjlist.begin(), v1.adjlist.end(),
				v2.adjlist.begin(), v2.adjlist.end(),
				std::inserter(nadj, nadj.begin()));
			hd_ymr = 0ul;
			/*
			 * Count all high degree adjacencies.
			 * NOTE: precisely, we should count the high degree adjacencies
			 * after coalescing the two vertexes. However, for simplicity,
			 * we use the CURRENT (without coalescing) degree instead.
			 * They may not be the same. For example, suppose that
			 * both A and B are adjacent with C, then before coalescing,
			 * degree of C is 2, but after coalescing, the degree is 1!
			 */
			for (auto & a : nadj)
				if (vertexes_.at(a).is_float == v1.is_float &&
					degree_of(a) >= (v1.is_float ? NAFPR : NAGPR))
					++hd_ymr;
			if (hd_ymr < (v1.is_float ? NAFPR : NAGPR))
			{
				flag_found = 1ul;
				break;
			}
		}
	
	if (flag_found != 0ul)
	{
		if (vertexes_.at(cpit->igv2).allocated)
		{/*
		 * If `igv2` is preallocated, we should coalesce `igv1` with `igv2`.
		 * So, we do an exchange here.
		 */
			if (vertexes_.at(cpit->igv1).allocated)
				throw std::logic_error("Both vertex " + std::to_string(cpit->igv2)
					+ " and " + std::to_string(cpit->igv1) + " are allocated");
			auto temp = cpit->igv1;
			cpit->igv1 = cpit->igv2;
			cpit->igv2 = temp;
		}
		// Do necessary works to coalesce `igv2` with `igv1`...
		vertexes_.at(cpit->igv2).valid = false;
		// Move all temps of `igv2` to `igv1`
		for (auto t_of_igv2 : vertexes_.at(cpit->igv2).temps)
			vertexes_.at(cpit->igv1).temps.insert(t_of_igv2);
		vertexes_.at(cpit->igv2).temps.clear();
		// Replace adjacencies in `igv2`'s adjacencies
		for (auto adj_of_igv2 : vertexes_.at(cpit->igv2).adjlist)
		{
			if (vertexes_.at(adj_of_igv2).adjlist.count(cpit->igv2) <= 0ul)
				throw std::logic_error("Adjacency " + std::to_string(adj_of_igv2)
					+ " of " + std::to_string(cpit->igv2) + " doesn't has it as adj");
			// It is safe to `erase()` for `std::set`, even if not existing
			vertexes_.at(adj_of_igv2).adjlist.erase(cpit->igv2);
			vertexes_.at(adj_of_igv2).adjlist.insert(cpit->igv1);
		}
		// Move all adjacencies of `igv2` to `igv1`
		for (auto adj_of_igv2 : vertexes_.at(cpit->igv2).adjlist)
			vertexes_.at(cpit->igv1).adjlist.insert(adj_of_igv2);
		vertexes_.at(cpit->igv2).adjlist.clear();
		// Deal with vertexes that have been removed
		for (auto vtx_rm : simplified_stack_)
			if (vertexes_.at(vtx_rm).adjlist.count(cpit->igv2) > 0ul)
			{
				vertexes_.at(vtx_rm).adjlist.erase(cpit->igv2);
				vertexes_.at(vtx_rm).adjlist.insert(cpit->igv1);
			}
		// Deal with other copies
		for (cpit2 = copies_.begin(); cpit2 != copies_.end(); ++cpit2)
			if (cpit2 != cpit)
			{
				size_t third_edge;
				if (cpit2->igv1 == cpit->igv2)
				{
					cpit2->igv1 = cpit->igv1;
					third_edge = search_copy(cpit2->igv2, cpit->igv1, false);
					if (third_edge != SIZE_MAX)
					{/*
					 * If 3 copy edges forms a triangle (`cpit`, `cpit2`, `third_edge`),
					 * then just invalidate the copy edge pointed by `cpit2`.
					 */
						for (auto inst_of_cp2 : cpit2->insts)
							copies_.at(third_edge).insts.insert(inst_of_cp2);
						cpit2->insts.clear();
						cpit2->status = HIGCopy::REDUNDANT;
					}
				}
				else if (cpit2->igv2 == cpit->igv2)
				{
					cpit2->igv2 = cpit->igv1;
					third_edge = search_copy(cpit2->igv1, cpit->igv1, false);
					if (third_edge != SIZE_MAX)
					{
						for (auto inst_of_cp2 : cpit2->insts)
							copies_.at(third_edge).insts.insert(inst_of_cp2);
						cpit2->insts.clear();
						cpit2->status = HIGCopy::REDUNDANT;
					}
				}
			}
		// Set the status
		cpit->status = HIGCopy::COALESCED;
	}
	return flag_found;
}

/*
 * When the graph cannot be simplified or coalesced,
 * try to find a copy edge with at least a low degree vertex.
 * Remove it to create a new low degree vertex by setting its
 * `status` as `FROZEN`.
 * Return non-zero integer on success.
 */
uint64_t HWorkingIGraph::freeze(void)
{
	uint64_t flag_found = 0ul;
	bool is_float;
	for (auto & c : copies_)
		if (c.status == HIGCopy::NORMAL)
		{
			is_float = vertexes_.at(c.igv1).is_float;
			if (is_float != vertexes_.at(c.igv2).is_float)
				throw std::logic_error("Valid copy edge found "
					"between float and integer " + std::to_string(c.igv1)
					+ ',' + std::to_string(c.igv2));
			if (degree_of(c.igv1) < (is_float ? NAFPR : NAGPR)
				|| degree_of(c.igv2) < (is_float ? NAFPR : NAGPR))
			{	// If copy edge with low vertex is found, then freeze it
				flag_found = 1ul;
				c.status = HIGCopy::FROZEN;
				break;
			}
		}
	return flag_found;
}

/*
 * When there are no vertexes to simplify, coalesce or freeze,
 * try to remove a high degree vertex and push it into the stack.
 */
uint64_t HWorkingIGraph::spill_potential(void)
{
	size_t i, n = vertexes_.size();
	std::vector<HIGVertex>::iterator vit;
	for (i = 0ul, vit = vertexes_.begin(); i < n; ++i, ++vit)
		if (vit->valid && vit->allocated == false)	// valid and not preallocated
			break;
	if (i >= n)
		return 0ul;
	else
	{
		simplified_stack_.push_back(i);
		vit->valid = false;
		for (auto u : vit->adjlist)
			vertexes_.at(u).adjlist.erase(i);
		/*
		 * When a vertex with valid copy edges is removed, we must
		 * freeze all the copy edges with it!
		 */
		size_t cidx;
		while ((cidx = search_copy(i)) != SIZE_MAX)
			copies_.at(cidx).status = HIGCopy::FROZEN;
		return 1ul;
	}
}

/*
 * When all vertexes (except preallocated) are removed and pushed
 * into stack, the selecting will begin.
 * Pop every vertex and try to allocate a register for it.
 * Return non-zere integer on success or 0 on failure.
 */
uint64_t HWorkingIGraph::select(void)
{
	size_t vidx;
	const RV64GFR *preg;
	size_t nregs;
	// The vertex at the top of stack
	std::vector<HIGVertex>::iterator cur_vtx;
	// Flag of the current register is conflicting
	bool flag_conflicting = false;

	while (simplified_stack_.empty() == false)
	{
		// Pop a vertex from the stack
		vidx = simplified_stack_.at(simplified_stack_.size() - 1ul);
		simplified_stack_.pop_back();
		cur_vtx = vertexes_.begin() + vidx;

		if (cur_vtx->is_float)
			preg = Available_FPR, nregs = NAFPR;
		else
			preg = Available_GPR, nregs = NAGPR;
		// Try all available registers
		flag_conflicting = true;
		while (nregs-- > 0ul)
		{
			flag_conflicting = false;
			for (auto adj : cur_vtx->adjlist)
			{
				if (vertexes_.at(adj).valid == false ||
					vertexes_.at(adj).allocated == false)
					throw std::logic_error("Adj " + std::to_string(adj) + " of "
						+ std::to_string(vidx) + " is invalid or not allocated");
				if (vertexes_.at(adj).reg == *preg)
				{
					flag_conflicting = true;
					break;
				}
			}
			if (flag_conflicting)
				++preg;
			else
				break;
		}
		
		if (flag_conflicting == false)
		{	// The vertex has been allocated a safe register!
			cur_vtx->valid = true;
			for (auto adj : cur_vtx->adjlist)
				vertexes_.at(adj).adjlist.insert(vidx);
			cur_vtx->allocated = true;
			cur_vtx->reg = *preg;
		}
		else
			// No available register for this vertex!
			break;
	}

	return flag_conflicting ? 0ul : 1ul;
}

/*
 * When the selection is failed, we must choose at least one temp
 * and give up allocating a register for it.
 * Strategy concerning degree and reference frequency will be used.
 * Temporaries in `already_spilled` will not be chosen again.
 * NOTE: This function is only available for base class and should
 * 	be deleted in derivied class!
 */
HIRTempno HInterferenceGraph::spill_actual(
	const std::set<HIRTempno> &already_spilled,
	const std::vector<HIRTemp> &temps) const
{
	size_t vidx, n = vertexes_.size(), min_vidx = SIZE_MAX;
	double min_cost = 0.0;
	HIRTempno t;

	// Try to find a valid vertex that is not in `already_spilled` with min cost
	for (vidx = 0ul; vidx < n; ++vidx)
	{
		if (vertexes_.at(vidx).valid &&
			already_spilled.count(t = *vertexes_.at(vidx).temps.begin()) == 0ul)
			if (min_cost == 0.0 ||
				temps.at(t).spillcost / double(degree_of(vidx)) < min_cost)
			{
				min_cost = temps.at(t).spillcost / double(degree_of(vidx));
				min_vidx = vidx;
			}
	}
	if (min_vidx == SIZE_MAX)
		// Oops! Not found...
		throw std::logic_error("Failed to choose spilled temp");
	return *vertexes_.at(min_vidx).temps.begin();
}

/*
 * Driver function of register allocation based on
 * interference graph...
 * Make sure that the original interference graph has been built!
 * The return value is reserved.
 */
uint64_t RV64Funct::reg_alloc(void)
{
#define PRINT_LOG(prefix)								\
	(opts.log_file != nullptr ? (						\
		log_output << prefix << std::endl,				\
		wintgraph_.print(log_output, '\0'),				\
		log_output << std::endl,						\
		true)											\
	: true)


	std::set<HIRTempno> spilled;
	bool flag_success = false;
	do {
		wintgraph_.build_from_oig(intgraph_, spilled);
		PRINT_LOG("New working interference graph:");
		do {
			do {
				do {
					do
						;
					while (wintgraph_.simplify() != 0ul);
					PRINT_LOG("After simplifying:");
				} while (wintgraph_.coalesce() != 0ul &&
				/*
				 * The macro `PRINT_LOG()` will always return `true`.
				 * So, if we use `&&` operator, the printing will happen
				 * only when the former function returns `true`, and
				 * it will not change the value of the condition!
				 */
					PRINT_LOG("After coalescing:"));
			} while (wintgraph_.freeze() != 0ul &&
				PRINT_LOG("After freezing:"));
		} while (wintgraph_.spill_potential() != 0ul &&
			PRINT_LOG("After spilling potential:"));
		// Now we can try to restore the graph and color it
		flag_success = wintgraph_.select();
		if (flag_success)
		{
			/*
			 * Scan the working interference graph and allocate
			 * register by setting reg desc and temp desc
			 * (for every valid vertex).
			 * If a callee-saved register is allocated, remember
			 * to add it in `saved_regs_` list!
			 */
			for (auto & vtx : wintgraph_.vertexes())
			{
				if (vtx.allocated)
				{	// Set descriptors if allocated
					for (auto t : vtx.temps)
					{
						/*
						 * Just setting these two descriptors are enough
						 * as `reset_desc()` will check them and set
						 * the descriptors more precisely.
						 */
						tdescs_.at(t).status = TempDesc::ALLOCATED;
						tdescs_.at(t).reg = vtx.reg;
						rdescs_[uint32_t(vtx.reg)].flag_allocated = true;
					}
					if (RV64GFR_is_callee_saved(vtx.reg))
						saved_regs_[vtx.reg] = INVALID_STLOC;
				}
			}
			PRINT_LOG("Allocation succeeded!");
		}
		else
		{
			// If failed, spill a temporary
			spilled.insert(intgraph_.spill_actual(spilled,
				ircode_->funct(hirfidx_).temps));
			if (opts.log_file != nullptr)
			{
				log_output << "Failed! ";
				for (auto st : spilled)
					log_output << '%' << st << ", ";
				log_output << "is(are) spilled." << std::endl;
			}
		}
	} while (flag_success == false);
	return 0ul;
#undef PRINT_LOG
}